package controllers

import (
	"bytes"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"image"
	"image/color"
	"image/draw"
	_ "image/gif"
	"image/jpeg"
	"image/png"
	"io"
	"net/http"
	"shudao-chat-go/models"
	"strings"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/s3"
	"github.com/beego/beego/v2/server/web"
)

type HazardController struct {
	web.Controller
}

type HazardRequest struct {
	//场景名称
	SceneName string `json:"scene_name"`
	//图片
	Image string `json:"image"`
	//用户ID
	UserID int `json:"user_id"`
}

// YOLOResponse 定义YOLO响应结构
type YOLOResponse struct {
	Boxes     [][]float64 `json:"boxes"`
	Labels    []string    `json:"labels"`
	Scores    []float64   `json:"scores"`
	ModelType string      `json:"model_type"`
}

// HazardResponse 定义隐患识别响应结构
type HazardResponse struct {
	Code int                    `json:"code"`
	Msg  string                 `json:"msg"`
	Data map[string]interface{} `json:"data"`
}

// 隐患识别
func (c *HazardController) Hazard() {
	var requestData HazardRequest
	if err := json.Unmarshal(c.Ctx.Input.RequestBody, &requestData); err != nil {
		c.Data["json"] = HazardResponse{
			Code: 400,
			Msg:  "参数错误",
		}
		c.ServeJSON()
		return
	}
	fmt.Println("requestData", requestData)

	// 验证参数
	if requestData.SceneName == "" || requestData.Image == "" {
		c.Data["json"] = HazardResponse{
			Code: 400,
			Msg:  "场景名称和图片链接不能为空",
		}
		c.ServeJSON()
		return
	}

	// 从OSS下载图片
	imageData, err := downloadImageFromOSS(requestData.Image)
	if err != nil {
		c.Data["json"] = HazardResponse{
			Code: 500,
			Msg:  "下载图片失败: " + err.Error(),
		}
		c.ServeJSON()
		return
	}

	// 验证图片数据
	if len(imageData) == 0 {
		c.Data["json"] = HazardResponse{
			Code: 500,
			Msg:  "下载的图片数据为空",
		}
		c.ServeJSON()
		return
	}

	fmt.Printf("下载图片成功，大小: %d 字节\n", len(imageData))

	// 将图片转换为base64
	imageBase64 := base64.StdEncoding.EncodeToString(imageData)

	// 获取YOLO API配置
	yoloBaseURL, err := web.AppConfig.String("yolo_base_url")
	if err != nil || yoloBaseURL == "" {
		yoloBaseURL = "http://172.16.35.50:18080" // 默认值
	}

	// 构建YOLO请求
	yoloRequest := map[string]interface{}{
		"modeltype":      requestData.SceneName,
		"image":          imageBase64,
		"conf_threshold": 0.5, // 默认置信度
	}

	jsonData, err := json.Marshal(yoloRequest)
	if err != nil {
		c.Data["json"] = HazardResponse{
			Code: 500,
			Msg:  "序列化请求失败: " + err.Error(),
		}
		c.ServeJSON()
		return
	}

	// 调用YOLO API
	resp, err := http.Post(
		fmt.Sprintf("%s/predict", yoloBaseURL),
		"application/json",
		bytes.NewBuffer(jsonData),
	)
	if err != nil {
		c.Data["json"] = HazardResponse{
			Code: 500,
			Msg:  "调用YOLO API失败: " + err.Error(),
		}
		c.ServeJSON()
		return
	}
	defer resp.Body.Close()

	// 读取响应
	body, err := io.ReadAll(resp.Body)
	if err != nil {
		c.Data["json"] = HazardResponse{
			Code: 500,
			Msg:  "读取YOLO响应失败: " + err.Error(),
		}
		c.ServeJSON()
		return
	}
	fmt.Println("body111", string(body))

	if resp.StatusCode != http.StatusOK {
		c.Data["json"] = HazardResponse{
			Code: resp.StatusCode,
			Msg:  "YOLO API返回错误: " + string(body),
		}
		c.ServeJSON()
		return
	}

	// 解析YOLO响应
	var yoloResp YOLOResponse
	if err := json.Unmarshal(body, &yoloResp); err != nil {
		c.Data["json"] = HazardResponse{
			Code: 500,
			Msg:  "解析YOLO响应失败: " + err.Error(),
		}
		c.ServeJSON()
		return
	}

	fmt.Println("yoloResp222", yoloResp)
	//如果labels为空，则返回错误
	if len(yoloResp.Labels) == 0 {
		c.Data["json"] = HazardResponse{
			Code: 500,
			Msg:  "没有识别到任何隐患",
		}
		c.ServeJSON()
		return
	}
	// 构建返回数据
	detectionResults := make([]map[string]interface{}, 0)
	for i, label := range yoloResp.Labels {
		if i < len(yoloResp.Scores) && i < len(yoloResp.Boxes) {
			result := map[string]interface{}{
				"label":   label,
				"score":   yoloResp.Scores[i],
				"box":     yoloResp.Boxes[i],
				"percent": fmt.Sprintf("%.2f%%", yoloResp.Scores[i]*100),
			}
			detectionResults = append(detectionResults, result)
		}
	}

	// 生成标注后的图片并上传到OSS
	var annotatedImageURL string
	if len(detectionResults) > 0 {
		// 解码原始图片
		fmt.Printf("开始解码图片，数据大小: %d 字节\n", len(imageData))

		// 检查图片格式
		img, format, err := image.Decode(bytes.NewReader(imageData))
		if err != nil {
			fmt.Printf("图片解码失败: %v\n", err)
			fmt.Printf("图片数据前20字节: %x\n", imageData[:min(20, len(imageData))])

			// 尝试手动检测PNG格式并转换
			if len(imageData) >= 8 && string(imageData[:8]) == "\x89PNG\r\n\x1a\n" {
				fmt.Printf("检测到PNG格式，尝试特殊处理...\n")

				// 尝试使用PNG解码器
				img, err = decodePNGImage(imageData)
				if err != nil {
					fmt.Printf("PNG特殊解码也失败: %v\n", err)
					c.Data["json"] = HazardResponse{
						Code: 500,
						Msg:  "PNG图片解码失败: " + err.Error() + "，请检查图片是否损坏",
					}
					c.ServeJSON()
					return
				}
				format = "png"
				fmt.Printf("PNG特殊解码成功\n")
			} else {
				c.Data["json"] = HazardResponse{
					Code: 500,
					Msg:  "解码图片失败: " + err.Error() + "，请检查图片格式是否支持（JPEG、PNG、GIF等）",
				}
				c.ServeJSON()
				return
			}
		}

		fmt.Printf("图片解码成功，格式: %s，尺寸: %dx%d\n", format, img.Bounds().Dx(), img.Bounds().Dy())

		// 绘制边界框
		annotatedImg := drawBoundingBox(img, yoloResp.Boxes, yoloResp.Labels, yoloResp.Scores)

		// 将标注后的图片编码为JPEG
		var buf bytes.Buffer
		if err := jpeg.Encode(&buf, annotatedImg, &jpeg.Options{Quality: 95}); err != nil {
			c.Data["json"] = HazardResponse{
				Code: 500,
				Msg:  "编码标注图片失败: " + err.Error(),
			}
			c.ServeJSON()
			return
		}

		// 生成文件名
		utcNow := time.Now().UTC()
		timestamp := utcNow.Unix()
		fileName := fmt.Sprintf("hazard_annotated/%d/%s_%d.jpg",
			utcNow.Year(),
			utcNow.Format("0102"),
			timestamp)

		// 上传标注后的图片到OSS
		annotatedImageURL, err = uploadImageToOSS(buf.Bytes(), fileName)
		if err != nil {
			c.Data["json"] = HazardResponse{
				Code: 500,
				Msg:  "上传标注图片到OSS失败: " + err.Error(),
			}
			c.ServeJSON()
			return
		}

		fmt.Printf("标注图片上传成功: %s\n", fileName)
	}

	// fmt.Println("annotatedImageURL", annotatedImageURL)
	if requestData.UserID == 0 {
		requestData.UserID = 1
	}

	//查询场景名称
	var scene models.Scene
	models.DB.Where("scene_en_name = ? and is_deleted = ?", requestData.SceneName, 0).First(&scene)
	if scene.ID == 0 {
		c.Data["json"] = map[string]interface{}{
			"statusCode": 500,
			"msg":        "场景名称不存在",
		}
		c.ServeJSON()
		return
	}

	var tx = models.DB.Begin()
	recognitionRecord := &models.RecognitionRecord{
		OriginalImageUrl:    requestData.Image,
		RecognitionImageUrl: annotatedImageURL,
		UserID:              requestData.UserID,
		Labels:              strings.Join(yoloResp.Labels, ", "), // 存储YOLO识别的标签
		Title:               scene.SceneName,
		TagType:             requestData.SceneName,
		SecondScene:         "", // 默认空字符串
		ThirdScene:          "", // 默认空字符串
	}
	if err := tx.Create(recognitionRecord).Error; err != nil {
		tx.Rollback()
		c.Data["json"] = map[string]interface{}{
			"statusCode": 500,
			"msg":        "识别失败: " + err.Error(),
		}
		c.ServeJSON()
		return
	}

	//获取labels对应的二级场景和三级场景
	var thirdSceneNames []string
	for _, label := range yoloResp.Labels {
		var secondScene models.SecondScene
		models.DB.Where("second_scene_name = ? and is_deleted = ?", label, 0).First(&secondScene)
		if secondScene.ID != 0 {
			tx.Create(&models.RecognitionRecordSecondScene{
				RecognitionRecordID: int(recognitionRecord.ID),
				SecondSceneID:       int(secondScene.ID),
			})

			// 通过secondScene.ID查询ThirdScene
			var thirdScene []models.ThirdScene
			models.DB.Where("second_scene_id = ? and is_deleted = ?", secondScene.ID, 0).Find(&thirdScene)
			if len(thirdScene) > 0 {
				for _, thirdScene := range thirdScene {
					thirdSceneNames = append(thirdSceneNames, thirdScene.ThirdSceneName)
				}
			}
		}
	}

	//对thirdSceneNames去重
	thirdSceneNames = removeDuplicates(thirdSceneNames)

	// 将三级场景名称数组更新到recognitionRecord的Description
	if len(thirdSceneNames) > 0 {
		// 将数组转换为空格分隔的字符串
		description := strings.Join(thirdSceneNames, " ")
		tx.Model(recognitionRecord).Update("Description", description)
	}

	tx.Commit()
	fmt.Println("thirdSceneNames", thirdSceneNames)
	// 返回成功响应
	c.Data["json"] = HazardResponse{
		Code: 200,
		Msg:  "识别成功",
		Data: map[string]interface{}{
			"scene_name":       requestData.SceneName,
			"total_detections": len(detectionResults),
			"detections":       detectionResults,
			"model_type":       yoloResp.ModelType,
			"original_image":   requestData.Image,
			"annotated_image":  annotatedImageURL, //前端用这个预链接渲染
			"labels":           strings.Join(yoloResp.Labels, ", "),
			"third_scenes":     thirdSceneNames, //三级场景名称数组
		},
	}
	c.ServeJSON()
}

// min 返回两个整数中的较小值
func min(a, b int) int {
	if a < b {
		return a
	}
	return b
}

// decodePNGImage 专门解码PNG图片
func decodePNGImage(imageData []byte) (image.Image, error) {
	// 直接使用PNG解码器
	img, err := png.Decode(bytes.NewReader(imageData))
	if err != nil {
		return nil, fmt.Errorf("PNG解码失败: %v", err)
	}
	return img, nil
}

// isImageContentType 检查Content-Type是否是图片类型
func isImageContentType(contentType string) bool {
	imageTypes := []string{
		"image/jpeg",
		"image/jpg",
		"image/png",
		"image/gif",
		"image/bmp",
		"image/webp",
		"image/tiff",
	}

	for _, imgType := range imageTypes {
		if contentType == imgType {
			return true
		}
	}
	return false
}

// downloadImageFromOSS 从OSS链接下载图片
func downloadImageFromOSS(imageURL string) ([]byte, error) {
	fmt.Printf("开始下载图片: %s\n", imageURL)

	client := &http.Client{
		Timeout: 30 * time.Second,
	}

	resp, err := client.Get(imageURL)
	if err != nil {
		return nil, fmt.Errorf("下载图片失败: %v", err)
	}
	defer resp.Body.Close()

	fmt.Printf("下载响应状态码: %d\n", resp.StatusCode)
	fmt.Printf("响应头: %+v\n", resp.Header)

	if resp.StatusCode != http.StatusOK {
		return nil, fmt.Errorf("下载图片失败，状态码: %d", resp.StatusCode)
	}

	imageData, err := io.ReadAll(resp.Body)
	if err != nil {
		return nil, fmt.Errorf("读取图片数据失败: %v", err)
	}

	fmt.Printf("图片下载完成，大小: %d 字节\n", len(imageData))

	// 检查Content-Type
	contentType := resp.Header.Get("Content-Type")
	if contentType != "" {
		fmt.Printf("图片Content-Type: %s\n", contentType)
		// 检查是否是图片类型
		if !isImageContentType(contentType) {
			fmt.Printf("警告: Content-Type不是图片类型: %s\n", contentType)
		}
	}

	return imageData, nil
}

// drawBoundingBox 在图像上绘制边界框和标签
func drawBoundingBox(img image.Image, boxes [][]float64, labels []string, scores []float64) image.Image {
	// 创建可绘制的图像副本
	bounds := img.Bounds()
	drawableImg := image.NewRGBA(bounds)
	draw.Draw(drawableImg, bounds, img, image.Point{}, draw.Src)

	// 红色边界框
	red := image.NewUniform(color.RGBA{255, 0, 0, 255})

	for _, box := range boxes {
		if len(box) >= 4 {
			x1, y1, x2, y2 := int(box[0]), int(box[1]), int(box[2]), int(box[3])

			// 绘制边界框
			drawRect(drawableImg, x1, y1, x2, y2, red)
		}
	}

	return drawableImg
}

// drawRect 绘制矩形
func drawRect(img *image.RGBA, x1, y1, x2, y2 int, color *image.Uniform) {
	bounds := img.Bounds()

	// 确保坐标在图像边界内
	if x1 < bounds.Min.X {
		x1 = bounds.Min.X
	}
	if y1 < bounds.Min.Y {
		y1 = bounds.Min.Y
	}
	if x2 >= bounds.Max.X {
		x2 = bounds.Max.X - 1
	}
	if y2 >= bounds.Max.Y {
		y2 = bounds.Max.Y - 1
	}

	// 设置线条粗细（像素数）
	lineThickness := 6

	// 绘制上边（粗线）
	for i := 0; i < lineThickness; i++ {
		y := y1 + i
		if y >= bounds.Max.Y {
			break
		}
		for x := x1; x <= x2; x++ {
			img.Set(x, y, color)
		}
	}

	// 绘制下边（粗线）
	for i := 0; i < lineThickness; i++ {
		y := y2 - i
		if y < bounds.Min.Y {
			break
		}
		for x := x1; x <= x2; x++ {
			img.Set(x, y, color)
		}
	}

	// 绘制左边（粗线）
	for i := 0; i < lineThickness; i++ {
		x := x1 + i
		if x >= bounds.Max.X {
			break
		}
		for y := y1; y <= y2; y++ {
			img.Set(x, y, color)
		}
	}

	// 绘制右边（粗线）
	for i := 0; i < lineThickness; i++ {
		x := x2 - i
		if x < bounds.Min.X {
			break
		}
		for y := y1; y <= y2; y++ {
			img.Set(x, y, color)
		}
	}
}

// removeDuplicates 去除字符串数组中的重复元素
func removeDuplicates(strs []string) []string {
	keys := make(map[string]bool)
	var result []string

	for _, str := range strs {
		if !keys[str] {
			keys[str] = true
			result = append(result, str)
		}
	}

	return result
}

// 复用shudaooss.go中的OSS配置变量

// uploadImageToOSS 上传图片数据到OSS并返回预签名URL
func uploadImageToOSS(imageData []byte, fileName string) (string, error) {
	// 获取S3会话
	s3Config := &aws.Config{
		Credentials:      credentials.NewStaticCredentials(ossAccessKey, ossSecretKey, ""),
		Endpoint:         aws.String(ossEndpoint),
		Region:           aws.String(ossRegion),
		S3ForcePathStyle: aws.Bool(true),
	}

	sess, err := session.NewSession(s3Config)
	if err != nil {
		return "", fmt.Errorf("创建S3会话失败: %v", err)
	}

	// 创建S3服务
	s3Client := s3.New(sess)

	// 上传图片到S3
	_, err = s3Client.PutObject(&s3.PutObjectInput{
		Bucket:      aws.String(ossBucket),
		Key:         aws.String(fileName),
		Body:        aws.ReadSeekCloser(strings.NewReader(string(imageData))),
		ContentType: aws.String("image/jpeg"),
		ACL:         aws.String("public-read"),
	})

	if err != nil {
		return "", fmt.Errorf("上传图片到OSS失败: %v", err)
	}

	// // 生成预签名URL（24小时有效期）
	// req, _ := s3Client.GetObjectRequest(&s3.GetObjectInput{
	// 	Bucket: aws.String(ossBucket),
	// 	Key:    aws.String(fileName),
	// })

	// presignedURL, err := req.Presign(24 * time.Hour)
	presignedURL := fmt.Sprintf("%s/%s/%s", ossEndpoint, ossBucket, fileName)
	if err != nil {
		return "", fmt.Errorf("生成预签名URL失败: %v", err)
	}

	return presignedURL, nil
}

// 前端传值步骤、json文件、封面图过来
type SaveStepRequest struct {
	AIConversationID uint64 `json:"ai_conversation_id"`
	Step             int    `json:"step"`
	PPTJsonUrl       string `json:"ppt_json_url"`
	CoverImage       string `json:"cover_image"`
	PPTJsonContent   string `json:"ppt_json_content"`
}

func (c *HazardController) SaveStep() {
	var requestData SaveStepRequest
	if err := json.Unmarshal(c.Ctx.Input.RequestBody, &requestData); err != nil {
		c.Data["json"] = map[string]interface{}{
			"statusCode": 400,
			"msg":        "请求数据解析失败",
		}
		c.ServeJSON()
		return
	}
	tx := models.DB.Begin()
	//更新到ai_conversation表
	if err := tx.Model(&models.AIConversation{}).Where("id = ?", requestData.AIConversationID).Updates(map[string]interface{}{
		"step":             requestData.Step,
		"ppt_json_url":     requestData.PPTJsonUrl,
		"cover_image":      requestData.CoverImage,
		"ppt_json_content": requestData.PPTJsonContent,
	}).Error; err != nil {
		tx.Rollback()
		c.Data["json"] = map[string]interface{}{
			"statusCode": 500,
			"msg":        "更新步骤失败",
		}
		c.ServeJSON()
		return
	}
	tx.Commit()
	c.Data["json"] = map[string]interface{}{
		"statusCode": 200,
		"msg":        "success",
	}
	c.ServeJSON()
}
